import torch
from torch.utils.data import Dataset, DataLoader
from torchgen import model
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification, DistilBertTokenizerFast
from torch.optim import AdamW
from sklearn.preprocessing import LabelEncoder
import json

# 加载模型
model.load_state_dict(torch.load('../models/emergency_distilbert.pt'))
model.eval()

# 示例预测
text = "我爸爸突然胸口疼，冒冷汗"
tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-multilingual-cased')
inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)

with torch.no_grad():
    outputs = model(
        input_ids=inputs['input_ids'],
        attention_mask=inputs['attention_mask']
    )

# 解析结果
# 编码标签
intent_encoder = LabelEncoder()
intent_encoder.fit(list(all_intents))
ner_tag_encoder = LabelEncoder()
ner_tag_encoder.fit(list(all_ner_tags))

intent_pred = intent_encoder.inverse_transform([torch.argmax(outputs['intent_logits']).item()])[0]
ner_preds = torch.argmax(outputs['ner_logits'], dim=-1).squeeze().tolist()
ner_tags = ner_tag_encoder.inverse_transform(ner_preds)

print(f"意图: {intent_pred}")
print("实体识别结果:")
for token, tag in zip(tokenizer.tokenize(text), ner_tags[:len(tokenizer.tokenize(text))]):
    print(f"{token}: {tag}")